#include "sphereface.h"
#include<vector>
using cv::Mat;
using std::vector;
namespace fr
{
	SphereFace::SphereFace(int batch_size,
		int channels,
		int height,
		int width,
		bool do_flip) :batch_size_(batch_size), channels_(channels), height_(height), width_(width), do_flip_(do_flip)
	{
		if (do_flip_)
		{
			batch_size_ = 2 * batch_size;
		}
	}

	vector<float> SphereFace::extract_feature(cv::Mat& aligned_face) const {
		vector<float> input_blob;
		preprocess(aligned_face, {}, input_blob);
		sphere20_->forward_pass(input_blob.data());
		return sphere20_->extract_feature();

	}

	void SphereFace::preprocess(const cv::Mat& src,
		const std::vector<cv::Point2f>& landmarks,
		std::vector<float>&blob_data) const {

		//resize
		Mat image = src.clone();
		image.convertTo(image, CV_32FC3, 1.0 / 128, -127.5 / 128);

		//NHWC to NCHW
		std::vector<cv::Mat> channels(3);
		cv::split(image, channels);
		for (auto &c : channels) {
			blob_data.insert(blob_data.end(), (float *)c.datastart, (float *)c.dataend);
		}


		if (do_flip_) {
			Mat image_flip;
			flip(image, image_flip, 0);
			std::vector<cv::Mat> channels_flip(3);
			cv::split(image_flip, channels_flip);
			for (auto &c : channels_flip) {
				blob_data.insert(blob_data.end(), (float *)c.datastart, (float *)c.dataend);
			}
		}
	}

	bool SphereFace::load_model(const std::string& init_nets,
		const std::string& predict_nets) {

		sphere20_ = std::make_unique<fr::NetCaffe>(init_nets, predict_nets, "fc5");
		return sphere20_->init(batch_size_, channels_, height_, width_);
	}

}


